import os
import shutil
import random

# 原始文件路径 （替换成你自己的路径）
image_folder = "/media/liuhw/38185FB1185F6CBE/disk_d/TC/3/JCXDXXJ/ZC/JPEGImages"  # PNG 图片所在文件夹
label_folder = "/media/liuhw/38185FB1185F6CBE/disk_d/TC/3/JCXDXXJ/ZC/txt"  # JSON 转换后的 YOLO txt 标签所在文件夹

# 目标数据集结构
dataset_root = "/media/liuhw/38185FB1185F6CBE/disk_d/TC/3/JCXDXXJ/ZC/dataset"  # 生成的根文件夹dataset 替换成你自己的路径
image_train_dir = os.path.join(dataset_root, "images/train")
image_val_dir = os.path.join(dataset_root, "images/val")
label_train_dir = os.path.join(dataset_root, "labels/train")
label_val_dir = os.path.join(dataset_root, "labels/val")

# 创建数据集目录
for dir_path in [image_train_dir, image_val_dir, label_train_dir, label_val_dir]:
    os.makedirs(dir_path, exist_ok=True)

# 获取所有图片名称（假设 PNG 文件和 TXT 标签名称一致）
image_files = [f for f in os.listdir(image_folder) if f.endswith(".jpg")]
random.shuffle(image_files)  # 打乱数据

# 数据集划分比例
train_ratio = 0.90
num_train = int(len(image_files) * train_ratio)

# 划分数据集
train_files = image_files[:num_train]
val_files = image_files[num_train:]

# 复制文件到对应目录
for file_list, img_dest, lbl_dest in [(train_files, image_train_dir, label_train_dir),
                                      (val_files, image_val_dir, label_val_dir)]:
    for img_file in file_list:
        # 复制图片
        shutil.copy(os.path.join(image_folder, img_file), os.path.join(img_dest, img_file))

        # 复制对应的标签
        label_file = img_file.replace(".jpg", ".txt")
        label_src = os.path.join(label_folder, label_file)
        if os.path.exists(label_src):  # 确保标签文件存在
            shutil.copy(label_src, os.path.join(lbl_dest, label_file))

print("数据集划分完成！")
